{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f0a8c09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import collections\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "694c35fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WorldEnv:\n",
    "    def __init__(self):\n",
    "        '''\n",
    "        环境初始化\n",
    "        功能：设置环境参数\n",
    "        参数说明：\n",
    "            distance_threshold：当智能体与目标的距离小于此值时，视为成功到达\n",
    "            action_bound：动作值的范围限制（动作被裁剪到[-1, 1]之间）\n",
    "        '''\n",
    "        self.distance_threshold = 0.15\n",
    "        self.action_bound = 1\n",
    "    \n",
    "    def reset(self):  # 重置环境\n",
    "        # 生成一个目标状态, 坐标范围是[3.5～4.5, 3.5～4.5]\n",
    "        \n",
    "        # 目标生成：在右上角区域（x∈[3.5,4.5], y∈[3.5,4.5]）随机生成目标位置\n",
    "        self.goal = np.array(\n",
    "            [4 + random.uniform(-0.5, 0.5), 4 + random.uniform(-0.5, 0.5)]\n",
    "        )\n",
    "        # 状态初始化：将智能体置于原点(0,0)\n",
    "        self.state = np.array([0, 0])\n",
    "        # 步数重置：计数器清零\n",
    "        self.count = 0\n",
    "        # 状态拼接：返回包含当前位置和目标位置的复合状态[x, y, gx, gy]\n",
    "        return np.hstack((self.state, self.goal))\n",
    "\n",
    "    def step(self, action):\n",
    "        # 动作裁剪：确保动作在[-1, 1]范围内，防止非法动作\n",
    "        action = np.clip(action, -self.action_bound, self.action_bound)\n",
    "        # 根据当前位置和动作计算新位置，使用max/min确保位置在环境边界[0, 5]×[0, 5]内\n",
    "        x = max(0, min(5, self.state[0] + action[0]))\n",
    "        y = max(0, min(5, self.state[1] + action[1]))\n",
    "        self.state = np.array([x, y])\n",
    "        self.count += 1\n",
    "        # 距离计算：计算当前位置与目标的欧氏距离\n",
    "        dis = np.sqrt(np.sum(np.square(self.state - self.goal)))\n",
    "        # 稀疏奖励：只有到达目标附近（距离≤0.15）时奖励为 0，否则为 - 1\n",
    "        reward = -1.0 if dis > self.distance_threshold else 0\n",
    "        # 成功条件：到达目标附近（距离≤0.15）\n",
    "        # 失败条件：达到最大步数（50 步）仍未到达\n",
    "        if dis <= self.distance_threshold or self.count == 50:\n",
    "            done = True\n",
    "        else:\n",
    "            done = False\n",
    "        # 拼接后的状态[新位置, 目标]， 奖励值， 终止标志\n",
    "        return np.hstack((self.state, self.goal)), reward, done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "451711ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolicyNet(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):\n",
    "        super().__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc3 = torch.nn.Linear(hidden_dim, action_dim)\n",
    "        # action_bound是环境可以接受的动作最大值\n",
    "        self.action_bound = action_bound\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc2(F.relu(self.fc1(x))))\n",
    "        return torch.tanh(self.fc3(x)) * self.action_bound\n",
    "\n",
    "class QValueNet(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim + action_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc3 = torch.nn.Linear(hidden_dim, 1)\n",
    "\n",
    "    def forward(self, x, a):\n",
    "        cat = torch.cat([x, a], dim = 1) # 拼接状态和动作\n",
    "        x = F.relu(self.fc2(F.relu(self.fc1(cat))))\n",
    "        return self.fc3(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "98447808",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DDPG:\n",
    "    ''' DDPG算法 '''\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,\n",
    "                 actor_lr, critic_lr, sigma, tau, gamma, device):\n",
    "        self.action_dim = action_dim\n",
    "        self.actor = PolicyNet(state_dim, hidden_dim, action_dim,\n",
    "                               action_bound).to(device)\n",
    "        self.critic = QValueNet(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.target_actor = PolicyNet(state_dim, hidden_dim, action_dim,\n",
    "                                      action_bound).to(device)\n",
    "        self.target_critic = QValueNet(state_dim, hidden_dim,\n",
    "                                       action_dim).to(device)\n",
    "        # 初始化目标价值网络并使其参数和价值网络一样\n",
    "        self.target_critic.load_state_dict(self.critic.state_dict())\n",
    "        # 初始化目标策略网络并使其参数和策略网络一样\n",
    "        self.target_actor.load_state_dict(self.actor.state_dict())\n",
    "        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),\n",
    "                                                lr=actor_lr)\n",
    "        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),\n",
    "                                                 lr=critic_lr)\n",
    "        self.gamma = gamma\n",
    "        self.sigma = sigma  # 高斯噪声的标准差,均值直接设为0\n",
    "        self.tau = tau  # 目标网络软更新参数\n",
    "        self.action_bound = action_bound\n",
    "        self.device = device\n",
    "\n",
    "    def take_action(self, state):\n",
    "        state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
    "        action = self.actor(state).detach().cpu().numpy()[0]\n",
    "        # 给动作添加噪声，增加探索\n",
    "        action = action + self.sigma * np.random.randn(self.action_dim)\n",
    "        return action\n",
    "\n",
    "    def soft_update(self, net, target_net):\n",
    "        for param_target, param in zip(target_net.parameters(),\n",
    "                                       net.parameters()):\n",
    "            param_target.data.copy_(param_target.data * (1.0 - self.tau) +\n",
    "                                    param.data * self.tau)\n",
    "\n",
    "    def update(self, transition_dict):\n",
    "        states = torch.tensor(transition_dict['states'],\n",
    "                              dtype=torch.float).to(self.device)\n",
    "        actions = torch.tensor(transition_dict['actions'],\n",
    "                               dtype=torch.float).to(self.device)\n",
    "        rewards = torch.tensor(transition_dict['rewards'],\n",
    "                               dtype=torch.float).view(-1, 1).to(self.device)\n",
    "        next_states = torch.tensor(transition_dict['next_states'],\n",
    "                                   dtype=torch.float).to(self.device)\n",
    "        dones = torch.tensor(transition_dict['dones'],\n",
    "                             dtype=torch.float).view(-1, 1).to(self.device)\n",
    "\n",
    "        next_q_values = self.target_critic(next_states,\n",
    "                                           self.target_actor(next_states))\n",
    "        q_targets = rewards + self.gamma * next_q_values * (1 - dones)\n",
    "        # MSE损失函数\n",
    "        critic_loss = torch.mean(\n",
    "            F.mse_loss(self.critic(states, actions), q_targets))\n",
    "        self.critic_optimizer.zero_grad()\n",
    "        critic_loss.backward()\n",
    "        self.critic_optimizer.step()\n",
    "\n",
    "        # 策略网络就是为了使Q值最大化\n",
    "        actor_loss = -torch.mean(self.critic(states, self.actor(states)))\n",
    "        self.actor_optimizer.zero_grad()\n",
    "        actor_loss.backward()\n",
    "        self.actor_optimizer.step()\n",
    "\n",
    "        self.soft_update(self.actor, self.target_actor)  # 软更新策略网络\n",
    "        self.soft_update(self.critic, self.target_critic)  # 软更新价值网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d607a699",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trajectory:\n",
    "    ''' 用来记录一条完整轨迹 '''\n",
    "    def __init__(self, init_state):\n",
    "        '''\n",
    "        初始化方法：\n",
    "            init_state：轨迹的起始状态，存入states列表\n",
    "            states：存储状态序列，初始化为包含初始状态的列表\n",
    "            actions：存储执行的动作序列\n",
    "            rewards：存储获得的奖励序列\n",
    "            dones：存储终止标志序列\n",
    "            length：记录轨迹长度（步数）\n",
    "        '''\n",
    "        self.states = [init_state]\n",
    "        self.actions = []\n",
    "        self.rewards = []\n",
    "        self.dones = []\n",
    "        self.length = 0\n",
    "\n",
    "    def store_step(self, action, state, reward, done):\n",
    "        '''\n",
    "        store_step 方法：\n",
    "            存储一步交互数据：动作、下一个状态、奖励和终止标志\n",
    "            每次调用后轨迹长度加 1\n",
    "        '''\n",
    "        self.actions.append(action)\n",
    "        self.states.append(state)\n",
    "        self.rewards.append(reward)\n",
    "        self.dones.append(done)\n",
    "        self.length += 1\n",
    "\n",
    "class ReplayBuffer_Trajectory:\n",
    "    # 类功能：管理多个轨迹的经验回放缓冲区\n",
    "    ''' 存储轨迹的经验回放池 '''\n",
    "    def __init__(self, capacity):\n",
    "        '''\n",
    "        初始化方法：\n",
    "            capacity：缓冲区最大容量\n",
    "            buffer：使用deque存储轨迹，自动淘汰最早的轨迹\n",
    "        '''\n",
    "        self.buffer = collections.deque(maxlen=capacity)\n",
    "    \n",
    "    def add_trajectory(self, trajectory):\n",
    "        '''\n",
    "        add_trajectory 方法：\n",
    "            将完整轨迹添加到缓冲区\n",
    "        '''\n",
    "        self.buffer.append(trajectory)\n",
    "\n",
    "    def size(self):\n",
    "        return len(self.buffer)\n",
    "    \n",
    "    def sample(self, batch_size, use_her, dis_threshold=0.15, her_ratio=0.8):\n",
    "        '''\n",
    "        sample方法\n",
    "            方法参数：\n",
    "                batch_size：采样的样本数\n",
    "                use_her：是否使用 HER 算法\n",
    "                dis_threshold：距离阈值，判断是否达到目标\n",
    "                her_ratio：使用 HER 的样本比例\n",
    "            采样过程：\n",
    "                随机选择一条轨迹\n",
    "                随机选择轨迹中的一个时间步\n",
    "                提取该时间步的状态、动作、奖励等信息\n",
    "        '''\n",
    "        batch = dict(states=[], actions=[], next_states=[], rewards=[], dones=[])\n",
    "        for _ in range(batch_size):\n",
    "            # 随机选择一条轨迹\n",
    "            traj = random.sample(self.buffer, 1)[0]\n",
    "            # 随机选择轨迹中的一个时间步\n",
    "            step_state = np.random.randint(traj.length)\n",
    "            # 提取该时间步的状态、动作、奖励等信息\n",
    "            state = traj.states[step_state]\n",
    "            next_state = traj.states[step_state + 1]\n",
    "            action = traj.actions[step_state]\n",
    "            reward = traj.rewards[step_state]\n",
    "            done = traj.dones[step_state]\n",
    "\n",
    "            '''\n",
    "            state.shape (4,)\n",
    "            next_state.shape (4,)\n",
    "            action.shape (2,)\n",
    "            '''\n",
    "\n",
    "            if use_her and np.random.uniform() <= her_ratio:\n",
    "                '''\n",
    "                HER 算法实现部分 算法逻辑：\n",
    "                    目标选择：从当前时间步之后的状态中随机选择一个作为新目标\n",
    "                    距离计算：计算下一状态与新目标的欧氏距离\n",
    "                    奖励重塑：根据距离重新计算奖励（未达到目标为 - 1，达到为 0）\n",
    "                    终止标志更新：根据距离更新终止标志\n",
    "                    状态重构：将原始状态的前两维与新目标拼接，形成新状态表示\n",
    "                '''\n",
    "                # 目标选择：从当前时间步之后的状态中随机选择一个作为新目标\n",
    "                step_goal = np.random.randint(step_state + 1, traj.length + 1)\n",
    "                # 获取目标的坐标。states的前两维，x y\n",
    "                goal = traj.states[step_goal][:2] # 使用HER算法的future方案设置目标\n",
    "                '''\n",
    "                在强化学习的上下文中，state[:2]通常表示状态向量的前两个维度，它们一般代表智能体的位置信息（如 x 坐标和 y 坐标）。\n",
    "                在基于目标的强化学习任务中，状态通常由两部分组成：智能体的当前状态和目标状态。通过将这两部分拼接，可以让智能体明\n",
    "                确知道自己当前的位置和需要到达的目标位置。\n",
    "                状态结构：\n",
    "                    原始状态state可能包含4个维度，例如[x, y, gx, gy]，其中前两位[x, y]表示智能体的位置\n",
    "                    目标状态goal通常也是位置信息，例如[gx, gy]\n",
    "                HER 算法的工作原理：\n",
    "                    HER 通过 \"事后诸葛\" 的方式，将失败的轨迹重新标记为成功轨迹\n",
    "                    具体做法是：从当前轨迹中选择一个后续状态作为新目标，然后重新计算奖励和终止条件\n",
    "                状态重构的逻辑：\n",
    "                    将原始状态的位置部分state[:2]与新目标goal拼接\n",
    "                    形成新的状态表示[当前位置, 目标位置]\n",
    "                    这样智能体就能学习到 \"从当前位置到达目标位置\" 的策略\n",
    "                '''\n",
    "                dis = np.sqrt(np.sum(np.square(next_state[:2] - goal)))\n",
    "                reward = -1.0 if dis > dis_threshold else 0\n",
    "                done = False if dis > dis_threshold else True\n",
    "                state = np.hstack((state[:2], goal))\n",
    "                next_state = np.hstack((next_state[:2], goal))\n",
    "            \n",
    "            batch['states'].append(state)\n",
    "            batch['next_states'].append(next_state)\n",
    "            batch['actions'].append(action)\n",
    "            batch['rewards'].append(reward)\n",
    "            batch['dones'].append(done)\n",
    "        \n",
    "        batch['states'] = np.array(batch['states'])\n",
    "        batch['next_states'] = np.array(batch['next_states'])\n",
    "        batch['actions'] = np.array(batch['actions'])\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "90794790",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iter 0:   0%|          | 0/200 [00:00<?, ?it/s]/tmp/ipykernel_987/3560096994.py:28: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /pytorch/torch/csrc/utils/tensor_new.cpp:254.)\n",
      "  state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
      "Iter 0: 100%|█████████▉| 199/200 [00:03<00:00, 64.51it/s, episode=190, return=-50.000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state.shape (4,)\n",
      "next_state.shape (4,)\n",
      "action.shape (2,)\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'float' object has no attribute 'shape'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[8]\u001b[39m\u001b[32m, line 41\u001b[39m\n\u001b[32m     39\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m replay_buffer.size() >= minimal_episodes:\n\u001b[32m     40\u001b[39m     \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(n_train):\n\u001b[32m---> \u001b[39m\u001b[32m41\u001b[39m         transition_dict = \u001b[43mreplay_buffer\u001b[49m\u001b[43m.\u001b[49m\u001b[43msample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[32m     42\u001b[39m         agent.update(transition_dict)\n\u001b[32m     43\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m (i_episode + \u001b[32m1\u001b[39m) % \u001b[32m10\u001b[39m == \u001b[32m0\u001b[39m:\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[7]\u001b[39m\u001b[32m, line 80\u001b[39m, in \u001b[36mReplayBuffer_Trajectory.sample\u001b[39m\u001b[34m(self, batch_size, use_her, dis_threshold, her_ratio)\u001b[39m\n\u001b[32m     78\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m'\u001b[39m\u001b[33mnext_state.shape\u001b[39m\u001b[33m'\u001b[39m, next_state.shape)\n\u001b[32m     79\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m'\u001b[39m\u001b[33maction.shape\u001b[39m\u001b[33m'\u001b[39m, action.shape)\n\u001b[32m---> \u001b[39m\u001b[32m80\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m'\u001b[39m\u001b[33mreward.shape\u001b[39m\u001b[33m'\u001b[39m, \u001b[43mreward\u001b[49m\u001b[43m.\u001b[49m\u001b[43mshape\u001b[49m)\n\u001b[32m     81\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m'\u001b[39m\u001b[33mdone.shape\u001b[39m\u001b[33m'\u001b[39m, done.shape)\n\u001b[32m     83\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m use_her \u001b[38;5;129;01mand\u001b[39;00m np.random.uniform() <= her_ratio:\n",
      "\u001b[31mAttributeError\u001b[39m: 'float' object has no attribute 'shape'"
     ]
    }
   ],
   "source": [
    "actor_lr = 1e-3\n",
    "critic_lr = 1e-3\n",
    "hidden_dim = 128\n",
    "state_dim = 4\n",
    "action_dim = 2\n",
    "action_bound = 1\n",
    "sigma = 0.1\n",
    "tau = 0.005\n",
    "gamma = 0.98\n",
    "num_episodes = 2000\n",
    "n_train = 20\n",
    "batch_size = 256\n",
    "minimal_episodes = 200\n",
    "buffer_size = 10000\n",
    "device = 'cpu'\n",
    "\n",
    "random.seed(0)\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "env = WorldEnv()\n",
    "replay_buffer = ReplayBuffer_Trajectory(buffer_size)\n",
    "agent = DDPG(state_dim, hidden_dim, action_dim, action_bound, actor_lr, critic_lr, sigma, tau, gamma, device)\n",
    "\n",
    "return_list = []\n",
    "for i in range(10):\n",
    "    with tqdm(total=int(num_episodes / 10), desc= 'Iter %d' % i) as pbar:\n",
    "        for i_episode in range(int(num_episodes / 10)):\n",
    "            episode_return = 0\n",
    "            state = env.reset()\n",
    "            traj = Trajectory(state)\n",
    "            done = False\n",
    "            while not done:\n",
    "                action = agent.take_action(state)\n",
    "                state, reward, done = env.step(action)\n",
    "                episode_return += reward\n",
    "                traj.store_step(action, state, reward, done)\n",
    "            replay_buffer.add_trajectory(traj)\n",
    "            return_list.append(episode_return)\n",
    "            if replay_buffer.size() >= minimal_episodes:\n",
    "                for _ in range(n_train):\n",
    "                    transition_dict = replay_buffer.sample(batch_size, True)\n",
    "                    agent.update(transition_dict)\n",
    "            if (i_episode + 1) % 10 == 0:\n",
    "                pbar.set_postfix(\n",
    "                    {'episode': '%d' % (num_episodes / 10 * i + i_episode + 1),\n",
    "                     'return': '%.3f' % np.mean(return_list[-10:])\n",
    "                     })\n",
    "            pbar.update(1)\n",
    "episodes_list = list(range(len(return_list)))\n",
    "plt.plot(episodes_list, return_list)\n",
    "plt.xlabel('Episodes')\n",
    "plt.ylabel('Returns')\n",
    "plt.title('DDPG with HER on {}'.format('GridWorld'))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
